import os
import shutil
from pathlib import Path

# 定义路径
validation_dir = "./results/results3d/validation"
test_dir = "../OpenLane/lane3d_1000/test"
output_dir1 = "./results/results3d/test"

# 仅保留curve_case场景
scene_txt_mapping = {
    "curve_case": "1000_curve.txt"
}

# 遍历每个场景（现在只有curve_case）
for scene, txt_file in scene_txt_mapping.items():
    txt_path = os.path.join(test_dir, txt_file)
    
    if not os.path.exists(txt_path):
        print(f"跳过 {scene}，未找到 {txt_path}")
        continue

    # 读取txt文件中的路径并清理
    with open(txt_path, 'r') as f:
        lines = [line.strip() for line in f.readlines() if line.strip()]
        lines = [line.replace("validation/", "") for line in lines]
        lines = [line.replace("\\", "/") for line in lines]

    # 为当前场景创建输出目录
    scene_output_dir = os.path.join(output_dir1, scene)
    os.makedirs(scene_output_dir, exist_ok=True)

    # 复制文件
    for rel_path in lines:
        # 分割路径获取segment目录名
        path_parts = rel_path.split('/')
        segment_dir = path_parts[0] if len(path_parts) > 1 else ""
        file_name = path_parts[-1]
        # 构建完整路径（确保不重复validation）
        src_path = os.path.join(validation_dir, rel_path)
        dst_path = os.path.join(scene_output_dir, rel_path)

        if not os.path.exists(src_path):
            if src_path.endswith('.jpg'):
                alt_src_path = src_path.replace('.jpg', '.json')
                if os.path.exists(alt_src_path):
                    src_path = alt_src_path
                    dst_path = dst_path.replace('.jpg', '.json')

        # 确保目标目录存在
        os.makedirs(os.path.dirname(dst_path), exist_ok=True)

        if os.path.exists(src_path):
            shutil.copy2(src_path, dst_path)
            print(f"已复制: {src_path} -> {dst_path}")
        else:
            print(f"警告: 源文件不存在 - {src_path} (请检查txt文件中的路径是否正确)")

print("curve_case场景文件处理完成！")